package com.sccba.uop.demo.util.algorithm;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.security.cert.CertificateException;
import javax.security.cert.X509Certificate;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Enumeration;

public class RSAUtils {

    public static byte[] readPublicKeyFromCer(String filePath) {
        try {
            X509Certificate certificate = X509Certificate.getInstance(Files.newInputStream(Paths.get(filePath)));
            return certificate.getPublicKey().getEncoded();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (CertificateException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] readPrivateKeyFromPfx(String filePath, String password) {
        try {
            KeyStore keyStore = KeyStore.getInstance("PKCS12");
            keyStore.load(Files.newInputStream(Paths.get(filePath)), password.toCharArray());
            Enumeration<String> aliases = keyStore.aliases();
            if (aliases.hasMoreElements()) {
                String keyAlias = aliases.nextElement();
                return keyStore.getKey(keyAlias, password.toCharArray()).getEncoded();
            }
        } catch (IOException e) {
            e.printStackTrace();
        } catch (KeyStoreException | NoSuchAlgorithmException | UnrecoverableKeyException | java.security.cert.CertificateException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] encrypt(byte[] pubKey, byte[] plainData)  {
        try {

            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.ENCRYPT_MODE, getPublicKey(pubKey));
            return cipher.doFinal(plainData);
        } catch (NoSuchPaddingException | IllegalBlockSizeException | NoSuchAlgorithmException | BadPaddingException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] decrypt(byte[] privKey, byte[] cipherData) {
        try {
            Cipher cipher = Cipher.getInstance("RSA");
            cipher.init(Cipher.DECRYPT_MODE, getPrivateKey(privKey));
            return cipher.doFinal(cipherData);
        } catch (NoSuchPaddingException | IllegalBlockSizeException | NoSuchAlgorithmException | BadPaddingException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static byte[] signWithSHA256(byte[] privKey, byte[] data) {
        try {
            Signature signature = Signature.getInstance("SHA256WithRSA");
            signature.initSign(getPrivateKey(privKey));
            signature.update(data);
            return signature.sign();
        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return null;
    }

    public static boolean verifyWithSHA256(byte[] pubKey, byte[] data, byte[] signatureData) {
        try {
            Signature signature = Signature.getInstance("SHA256WithRSA");
            signature.initVerify(getPublicKey(pubKey));
            signature.update(data);
            return signature.verify(signatureData);
        } catch (NoSuchAlgorithmException | SignatureException | InvalidKeyException e) {
            e.printStackTrace();
        }
        return false;
    }

    private static PublicKey getPublicKey(byte[] pubKey) {
        try {
            X509EncodedKeySpec keySpec = new X509EncodedKeySpec(pubKey);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PublicKey publicKey = keyFactory.generatePublic(keySpec);
            return publicKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
        }
        return null;
    }

    private static PrivateKey getPrivateKey(byte[] privKey) {
        try {
            PKCS8EncodedKeySpec keySpec = new PKCS8EncodedKeySpec(privKey);
            KeyFactory keyFactory = KeyFactory.getInstance("RSA");
            PrivateKey privateKey = keyFactory.generatePrivate(keySpec);
            return privateKey;
        } catch (NoSuchAlgorithmException | InvalidKeySpecException e) {
            e.printStackTrace();
        }
        return null;
    }

}
